//  原版权信息 LAC v3 

//  修改者:
//  Copyright (c) 2025 刘子瑶
//  所属单位: 中国科学院信息工程研究所网络空间安全防御重点实验室
//           Key Laboratory of Cyberspace Security Defense, Institute of Information Engineering, CAS
//           中国科学院大学网络空间安全学院
//           School of Cyber Security, University of Chinese Academy of Sciences     

//  modification:
//  (1) replaced BCH encoding with polar encoding in the function `encode_to_e2`
//  (2) replaced BCH decoding with polar decoding in the function `original_pke_dec`
//  (3) added functions `polar_enc_cpucycles`, `polar_dec_cpucycles` to evaluate the performance (CPU cycles) of polar code encoding/decoding


#include <string.h>
#include <stdio.h>
#include <immintrin.h>
#include <x86intrin.h>
#include "api.h"
#include "rand.h"
#include "ecc.h"
#include "bin-lwe.h"
#include "ntt.h"
#include "polarlac_param.h"
#include "fips202.h"

#define RATIO 126 // Q/2

#define MAX_RECORDS 100000 // number of tests
uint64_t polar_enc_cycle[MAX_RECORDS];
int polar_enc_index = 0;
uint64_t polar_dec_cycle[MAX_RECORDS];
int polar_dec_index = 0;


static int encode_to_e2(unsigned char *e2, const unsigned char *m, unsigned long long mlen, int *c2_len)
{
	int i;
	int8_t message;
	int vec_bound;
	unsigned char *p_code;

	/* polar encoding */
    uint8_t u[CODE_LEN*8]; // source sequence(each element stores 1 bits)
	uint8_t code[CODE_LEN]; // codeword sequence(each element stores 8 bits)
    memset(u, 0, sizeof(u));
	memset(code, 0, sizeof(code));

	// fill the message m into the source sequence
	int info_cnt = 0;
    for (int i = 0; i < CODE_LEN; i++)
    {
        for(int j = 0; j < 8; j++)
        {
            if(ecc_info_nodes[8*i+j] == 1)
            {
				u[8*i+j] = (uint8_t)(m[info_cnt/8] >> (info_cnt%8)) & 0x01;
                info_cnt++;
            }
        }
    }

	uint64_t polar_enc_start = __rdtsc();
	polar_ecc_enc(u);
	uint64_t polar_enc_end = __rdtsc();
	if (polar_enc_index < MAX_RECORDS) {
        polar_enc_cycle[polar_enc_index] = polar_enc_end - polar_enc_start;
		polar_enc_index++;
    }

	// each element stores 1 binary value -> each element stores 8 binary values
	for (int i = 0; i < CODE_LEN; i++)
	{
    	for (int j = 0; j < 8; j++) {
        	code[i] |= (u[i * 8 + j] << j);
    	}
	}

	p_code = (unsigned char *)code;
    
	
	//compute the length of c2
	*c2_len=CODE_LEN*8; // the code length of the ecc
	vec_bound=*c2_len;
	//compute code*q/2+e2
	for(i=0;i<vec_bound;i++)
	{
		//RATIO=q/2. add code*q/2 to e2
		message=RATIO*((p_code[i/8]>>(i%8))&1);
		e2[i]=e2[i]+message;
	}
	
	return 0;
}
/* END FOR POLAR */


//key generation
int crypto_encrypt_keypair( unsigned char *pk, unsigned char *sk)
{
	//check parameter
	if(pk==NULL || sk==NULL)
	{
		return -1;
	}
	kg(pk,sk);
	
	return 0;
}

//key generation with seed
int kg_seed(unsigned char *pk, unsigned char *sk, unsigned char *seed)
{
	unsigned char seeds[2*SEED_LEN];
	unsigned char a[DIM_N];
	unsigned char e[DIM_N];
	unsigned char h_pk[HASHLEN];
	//check pointer
	if(pk==NULL || sk==NULL)
	{
		return -1;
	}
	
	//generate two seeds for a,sk,e
	sha3_512(seeds, seed, SEED_LEN);
	memcpy(pk,seeds,SEED_LEN);
	
	//generate a
	gen_a(a,seeds);//print_bytes(sk,CRYPTO_SECRETKEYBYTES);

	keccak_state state;
    shake256_absorb_once(&state, seeds + SEED_LEN, SEED_LEN);
	//generate  sk,e
	gen_e(sk,&state);
	gen_e(e,&state);
	//compute pk=a*sk+e
	poly_aff(a,sk,e,pk+SEED_LEN,DIM_N);
	//copy pk=as+e to the second part of sk, now sk=s|pk
	memcpy(sk+SK_PART_LEN,pk,PK_LEN);
	sha3_256(h_pk,pk,PK_LEN);
	memcpy(sk+SK_PART_LEN+PK_LEN,h_pk,HASHLEN);
	
	return 0;
}

//key generation
int kg(unsigned char *pk, unsigned char *sk)
{
	unsigned char seed[SEED_LEN];
	
	//generate seed
	random_bytes(seed,SEED_LEN);		
	//key generation with seed 
	kg_seed(pk,sk,seed);	
	
	return 0;
}

//************************************ encryption ************************************//
int crypto_encrypt( unsigned char *c, unsigned long long *clen, const unsigned char *m, unsigned long long mlen, const unsigned char *pk, unsigned char *seeds)
{
	//check parameter
	if(c==NULL || m==NULL || pk==NULL)
	{
		return -1;
	}
	if(mlen>MESSAGE_LEN)
	{
		return -1;
	}
	
	//call pke encryption function
	original_pke_enc(pk,m, mlen,c,clen, seeds);

	return 0;
}

// encryption
int original_pke_enc(const unsigned char *pk, const unsigned char *m, unsigned long long mlen, unsigned char *c, unsigned long long *clen, unsigned char *seeds)
{
	original_pke_enc_seed(pk,m,mlen,c,clen,seeds);	

	return 0;
}

// encryption with seed
int original_pke_enc_seed(const unsigned char *pk, const unsigned char *m, unsigned long long mlen, unsigned char *c, unsigned long long *clen, unsigned char *seed)
{
	unsigned char r[DIM_N];
	unsigned char e1[DIM_N],e2[DIM_N];
	unsigned char c2[C2_VEC_NUM];
	unsigned char a[DIM_N];
	
	int c2_len;
	
	//check parameter
	if(pk==NULL || m==NULL || c==NULL )
	{
		return -1;
	}
	if(mlen>MESSAGE_LEN)
	{
		return -1;
	}

	//generate  a from seed in the first part of pk
	gen_a(a,pk);
	keccak_state state;
    shake256_absorb_once(&state, seed, SEED_LEN);
	//generate random vector r
	gen_e(r,&state);
	//generate error vector e1
	gen_e(e1,&state);
	//generate error vector e2
	gen_e(e2,&state);

	/* FOR POLAR */
	//encode message to e2
	encode_to_e2(e2,m,mlen,&c2_len);
	/* END FOR POLAR */

#if defined LAC_LIGHT
	unsigned char c1[DIM_N];
	//generate c1: c1=a*r+e1
	poly_aff(a,r,e1,c1,DIM_N);
	//compress c1
	poly_compress_c1_1bit(c1,c,DIM_N);

	//generate c2: c2=b*r+e2+m*[q/2]
	poly_aff(pk+SEED_LEN,r,e2,c2,c2_len);
	//compress c2
	poly_compress_c2_4bit(c2,c+DIM_N*7/8,c2_len);
	*clen=DIM_N*7/8+c2_len/2;
#elif defined LAC128
	unsigned char c1[DIM_N];
	//generate c1: c1=a*r+e1
	poly_aff(a,r,e1,c1,DIM_N);
	//compress c1
	poly_compress_c1_1bit(c1,c,DIM_N);

	//generate c2: c2=b*r+e2+m*[q/2]
	poly_aff(pk+SEED_LEN,r,e2,c2,c2_len);
	//compress c2
	poly_compress_c2_5bit(c2,c+DIM_N*7/8,c2_len);
	*clen=DIM_N*7/8+c2_len*3/8;
#else // LAC256
	//generate c1: c1=a*r+e1
	poly_aff(a,r,e1,c,DIM_N);

	//generate c2: c2=b*r+e2+m*[q/2]
	poly_aff(pk+SEED_LEN,r,e2,c2,c2_len);
	//compress c2
	poly_compress_c2_4bit(c2,c+DIM_N,c2_len);
	*clen=DIM_N+c2_len/2;
#endif	

	return 0;
}


//************************************ decryption ************************************//
int crypto_encrypt_open(unsigned char *m, unsigned long long *mlen,const unsigned char *c, unsigned long long clen,const unsigned char *sk)
{
	//check parameter
	if(sk==NULL || m==NULL || c==NULL || mlen==NULL)
	{
		return -1;
	}
	
	//call pke decryption function
	original_pke_dec(sk,c,clen,m,mlen);

	return 0;
}

// decrypt
int original_pke_dec(const unsigned char *sk, const unsigned char *c,unsigned long long clen, unsigned char *m, unsigned long long *mlen)
{
	//check parameter
	if(sk==NULL || m==NULL || c==NULL)
	{
		return -1;
	}
	
	unsigned char out[DIM_N];
	unsigned char c2[C2_VEC_NUM];	
	unsigned char m_buf[MESSAGE_LEN];
	float llr[CODE_LEN*8]; // log-likelihood ratio of the received signal
	
	int temp;
	int half = 126; // Q/2
	int half_2 = 63; // half/2
	uint8_t m_cap[MESSAGE_LEN*8]; // estimated message(each element stores 1 bit)
	uint64_t polar_dec_start = __rdtsc();

	uint64_t polar_dec_end = __rdtsc();
	
#if defined LAC_LIGHT
	unsigned char c1[DIM_N];
	int c2_len=(clen-DIM_N*7/8)*2;

	//c1 decompress
	poly_decompress_c1_1bit(c,c1,DIM_N);
	//c2 decompress
	poly_decompress_c2_4bit(c+DIM_N*7/8,c2,c2_len);
	//c1*sk
	poly_mul(c1,sk,out,c2_len);
#elif defined LAC128
	unsigned char c1[DIM_N];
	int c2_len=(clen-DIM_N*7/8)/3*8;

	//c1 decompress
	poly_decompress_c1_1bit(c,c1,DIM_N);
	//c2 decompress
	poly_decompress_c2_5bit(c+DIM_N*7/8,c2,c2_len);
	//c1*sk
	poly_mul(c1,sk,out,c2_len);
#else // LAC256
	int c2_len=(clen-DIM_N)*2;

	//c2 decompress
	poly_decompress_c2_4bit(c+DIM_N,c2,c2_len);
	//c1*sk
	poly_mul(c,sk,out,c2_len);
#endif

	/* FOR POLAR */
	*mlen=DATA_LEN;
	// compute llr
	for(int i = 0; i < c2_len; i++)
	{
		// compute m*q/2+e in [0,250]
		temp = (c2[i]-out[i]+Q)%Q;
		// m*q/2+e in [-63,187]=[-63,125]+[126,187]
		temp = temp - half_2;
		// m*q/2+e in [-125,125]
		if(temp >= half)
		{
			temp = temp - Q; // [126,187]——>[-125,-64]
		}
		llr[i] = -(float)temp / half_2; // 0 is modulated to -q/4, and 1 is modulated to q/4
	}

	//polar decode to recover m
	polar_ecc_dec(m_cap, llr);
	if (polar_dec_index < MAX_RECORDS) {
        polar_dec_cycle[polar_dec_index] = polar_dec_end - polar_dec_start;
		polar_dec_index++;
    }

	// each element stores 1 binary value -> each element stores 8 binary values
	memset(m_buf, 0, MESSAGE_LEN);
	for (int i = 0; i < MESSAGE_LEN; i++)
    {
        for(int j = 0; j < 8; j++)
        {
            m_buf[i] |= (m_cap[8*i+j] << j);
        }
    }

	memcpy(m, m_buf, *mlen);
	/* END FOR POLAR */
	
	return 0;
}

int polar_enc_cpucycles()
{
	uint64_t polar_enc_cycle_sum = 0;
	float polar_enc_cycle_avg;
    
	for(int i = 0; i < MAX_RECORDS; i++)
	{
		polar_enc_cycle_sum += polar_enc_cycle[i];
	}

	polar_enc_cycle_avg = (float)polar_enc_cycle_sum / MAX_RECORDS;

	printf("polar_enc  : %f cpucycles\n", polar_enc_cycle_avg);
	
	return 0;
}

int polar_dec_cpucycles()
{
	uint64_t polar_dec_cycle_sum = 0;
	float polar_dec_cycle_avg;
    
	for(int i = 0; i < MAX_RECORDS; i++)
	{
		polar_dec_cycle_sum += polar_dec_cycle[i];
	}

	polar_dec_cycle_avg = (float)polar_dec_cycle_sum / MAX_RECORDS;

	printf("polar_dec  : %f cpucycles\n", polar_dec_cycle_avg);
	
	return 0;
}